'''
该文件的第一个功能
COCO的验证集中存在着灰度图，需要将annotations中的灰度图像删除掉
000000007888.jpg
000000205289.jpg
000000209222.jpg
000000024021.jpg
000000061418.jpg
000000353180.jpg
000000431848.jpg
000000130465.jpg
000000141671.jpg
000000274219.jpg
这里我们采取的是手动删除内容的
'''
from collections import defaultdict
import json
from tools.check_RGBorGray import get_gray_jpg2
import os

def COCO_valid_clean_gray(val_img_path,s_val_json_path,d_val_json_path):
    '''
    将原始数据集中的Gray图像进行清理
    :param val_img_path:  验证集图片路径
    :param s_val_json_path: 待清除的val验证集文件
    :param d_val_json_path: 清除后的val验证集文件
    :return:
    '''
    # 在val_img_path路径下找到灰度图像
    gray_files = get_gray_jpg2(val_img_path)
    gray_files = list(map(lambda x: x.split('.')[0], gray_files))
    print('get gray images ok')
    # 读取s_val_json_path文件，并将与gray_files列表中文件相关的内容进行删除
    # 主要需要删除的内容为'annotations'和'images'两项
    images = list()
    annotations = list()
    with open(s_val_json_path, 'r') as f:
        json_data = json.load(f)
    for ann in json_data['annotations']:
        if str(ann['image_id']).zfill(12) not in gray_files:
            annotations.append(ann)
    for img in json_data['images']:
        if str(img['id']).zfill(12) not in gray_files:
            images.append(img)
    # 将修改后的json文件写入到d_val_json_path文件中
    json_data['annotations'] = annotations
    json_data['images'] = images
    with open(d_val_json_path, 'w') as f:
        f.write(json.dumps(json_data, indent=1, separators=(',', ':')))
#COCO_valid_clean_gray('../COSSY/coco/val2017','../COSSY/coco/annotations/instances_val2017.json',\
#                      '../COSSY/coco/annotations/instances_val2017_dropgray.json')
######################################################
'''
    COCO验证集数据格式处理，使得验证集的id为文件名，数据中仅包含人
    从训练集处理方式上入手，
    参照已有的标注数据，主要需要保住"annotations"和"images"
    并且"annotations"需要包含area，bbox，category_id，person_id，image_id，iscrowd，segmentation
    "images"需要包含"file_name": "Mar10_000002.jpg","id": "Mar10_000003","width"和"height"
    将segmentation清空
'''
def COCO_valid_set_format(s_val_json_path,d_val_json_path):
    coco = False
    imgid2anns = defaultdict(list)
    images = list()
    annotations = list()
    print(f'Loading annotations {s_val_json_path} into memory...')
    with open(s_val_json_path, 'r') as f:
        json_data = json.load(f)
    for ann in json_data['annotations']:
        ann['segmentation'] =[]      # 清空分割
        img_id = ann['image_id']
        # get width and height
        if len(ann['bbox']) == 4:
            # using COCO dataset. 4 = [x1,y1,w,h]
            coco = True
            # convert COCO format: x1,y1,w,h to x,y,w,h
            ann['bbox'][0] = ann['bbox'][0] + ann['bbox'][2] / 2
            ann['bbox'][1] = ann['bbox'][1] + ann['bbox'][3] / 2
            ann['bbox'].append(0)
            if ann['bbox'][2] > ann['bbox'][3]:
                ann['bbox'][2],ann['bbox'][3] = ann['bbox'][3],ann['bbox'][2]
                ann['bbox'][4] -= 90
        else:
            # using rotated bounding box datasets. 5 = [cx,cy,w,h,angle]
            assert len(ann['bbox']) == 5, 'Unknown bbox format' # x,y,w,h,a
        if ann['bbox'][4] == 90:               # 对数据集进行处理，根据周期性，将角度为90的变为角度为-90的
            ann['bbox'][4] = -90
        imgid2anns[img_id].append(ann)
    for img in json_data['images']:
        img_id = img['id']                    # 对img['id']进行处理
        anns = imgid2anns[img_id]
        # if there is crowd gt, skip this image
        if coco and any(ann['iscrowd'] for ann in anns):
            continue
        # only for person detection
        if not any(ann['category_id']==1 for ann in anns):
            continue
            # and ignore all other categories
        img['id'] = str(img['id']).zfill(12)     # 填充到12位
        images.append(img)        # 所有需要保留下来的图片
        imgid2anns[img_id] = [a for a in anns if a['category_id']==1]   # 将非人的目标过滤出去
        for a in imgid2anns[img_id]:
            a['image_id']=str(a['image_id']).zfill(12)
            annotations.append(a)
        # self.imgid2info[img['id']] = img
    json_data['annotations'] = annotations
    json_data['images'] = images
    # 将获取的xml内容写入到json文件中
    with open(d_val_json_path, 'w') as f:
        f.write(json.dumps(json_data, indent=1, separators=(',', ':')))
'''
    进行数据格式的转换
'''
#COCO_valid_set_format('../COCO/annotations/instances_val2017.json','../COCO/annotations/instances_val2017_new.json')

######################################################

'''
    进行json文件内容的校验，
    查看是否所有的包含人的json验证集图片有哪些
    以及这些图片的annotation中类别是否只有person这一类
'''
def check_new_valid_json(new_json_path):
    print(f'Checking annotations {new_json_path}...')
    with open(new_json_path, 'r') as f:
        json_data = json.load(f)
    print('length of effective images are:{}'.format(len(json_data['images'])))
    for ann in json_data['annotations']:
        if ann['category_id'] != 1:           # ann 类别标签不为1
            print(ann['image_id'])
#check_new_valid_json('../COCO/annotations/instances_val2017_new.json')
# 经过检测，有效的数据为2364个，并且这些图片的annotation中类别只有person这一类
#######################################################

def remove_val_imgdir_withoutPerson(img_dir,json_path):
    '''
    删除掉var_img_dir路径下的没有包含Person的图片
    :param img_dir: 验证集图片目录
    :param json_path: 经过COCO_valid_set_format之后的输出json文件
    :return:
    '''
    reserve_imgs = list()            # 需要保留下来的图片
    with open(json_path, 'r') as f:
        json_data = json.load(f)
    for img in json_data['images']:
        reserve_imgs.append(img['file_name'])
    print('get all reserved imgs ok')

    for file in os.listdir(img_dir):
        if file not in reserve_imgs:
            os.remove(os.path.join(img_dir,file))
    print('val_image dir clear ok!')
if __name__=='__main__':
    # 删除数据集中的灰度图得到instances_val2017_dropgray.json文件
    COCO_valid_clean_gray('../COSSY/coco/val2017', '../COSSY/coco/annotations/instances_val2017.json', \
                          '../COSSY/coco/annotations/instances_val2017_dropgray.json')
    COCO_valid_set_format('../COSSY/coco/annotations/instances_val2017_dropgray.json',
                          '../COSSY/coco/annotations/instances_val2017_setformat.json')
    '''
        经过测试，应该是没有问题的能够和我们手动得出一样的结果
    '''
    # 查看有多少个文件保留了下来
    #check_new_valid_json('../COSSY/coco/annotations/instances_val2017_setformat.json')
    # 清理掉COCO val集 图像中没有包含人的照片
    remove_val_imgdir_withoutPerson('../COSSY/coco/val2017',
                                    '../COSSY/coco/annotations/instances_val2017_setformat.json')


